#include "stdafx.h"
#include "PortCache.h"

PortCache::PortCache()
{
	RtlZeroMemory(_tcpPortTable, sizeof(_tcpPortTable));
	RtlZeroMemory(_udpPortTable, sizeof(_udpPortTable));
}

int PortCache::GetTcpPortPid(int port)
{
	if( _tcpPortTable[port] != 0 )
	{
		return _tcpPortTable[port];
	}
	else
	{
		// Clear the table
		RtlZeroMemory(_tcpPortTable, sizeof(_tcpPortTable));

		// Rebuild the table
		MIB_TCPTABLE_OWNER_PID table;
		table.dwNumEntries = sizeof(table) / sizeof(table.table[0]);

		DWORD tableSize = sizeof(table);

		if( GetExtendedTcpTable((void *)&table, &tableSize, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0) == NO_ERROR )
		{
			for(unsigned int i = 0; i < table.dwNumEntries; i++)
			{
				_tcpPortTable[ntohs((unsigned short)table.table[i].dwLocalPort)] = table.table[i].dwOwningPid;
			}
		}
		
		// Return
		return _tcpPortTable[port];
	}
}

int PortCache::GetUdpPortPid(int port)
{
	if( _udpPortTable[port] != 0 )
	{
		return _udpPortTable[port];
	}
	else
	{
		// Clear the table
		RtlZeroMemory(_udpPortTable, sizeof(_udpPortTable));

		// Rebuild the table
		MIB_UDPTABLE_OWNER_PID table;
		table.dwNumEntries = sizeof(table) / sizeof(table.table[0]);

		DWORD tableSize = sizeof(table);

		if( GetExtendedUdpTable((void *)&table, &tableSize, FALSE, AF_INET, UDP_TABLE_OWNER_PID, 0) == NO_ERROR)
		{
			for(unsigned int i = 0; i < table.dwNumEntries; i++)
			{
				_udpPortTable[ntohs((unsigned short)table.table[i].dwLocalPort)] = table.table[i].dwOwningPid;
			}
		}

		// Return
		return _udpPortTable[port];
	}
}